树相关函数的实现和ts封装

下面的内容关于封装的树的部分操作的工具类，平铺树转 tree，树转平铺、树的遍历、查找、筛选等。

树的自定义结构

treeCreate 函数接受一个平铺数组，其中平铺数组中的节点数据结构总是可变的，为了满足这个函数能够作为通用函数使用，在函数参数中提供了配置项，其中 id 和 faId 为树节点的 id 和父节点 id，children 为子节点的属性名。其中 ts 类型是比较完善的，你能够提供泛型去约束树节点类型，也可以不提供泛型，让其自动识别子节点类型，从而提供配置项中值的自动判断。

实现

enum TreeNodeTypeKeyEnum {
    id = "id",
    faId = "faId",
    children = "children",
}

/** 默认节点参数数据类型 */
export type TreeNodeItemType = {
    id: TreeNodeTypeKeyEnum.id;
    faId: TreeNodeTypeKeyEnum.faId;
    children: TreeNodeTypeKeyEnum.children;
};

/** 构建树函数配置项 */
type TreeCreateConfig<T extends Record<string, string>> = {
    keyMapping: Partial<Record<keyof Pick<TreeNodeItemType, "id" | "faId" | "children">, keyof T>>;
};

/** 检查是否是可索引类型 */
const checkIndexType = (data: any) => {
    return ["string", "number", "symbol"].includes(typeof data);
};

/** 树原型类型 */
type TreePrototypeType<T> = {
    treeForEach: (callback?: (item: T) => void) => ReturnType<typeof treeDFS<T>>;
    treeFind: (callback?: (item: T) => boolean) => ReturnType<typeof treeFind<T>>;
    treeFilter: (
        callback?: (item: T) => boolean,
        config?: {
            isSearch?: boolean;
        }
    ) => ReturnType<typeof treeFilter<T>>;
};

/** key为K类型中所有的key，以K[key]作为索引，用索引作为key取T中对应key的类型，该类型则为K[key]的类型 */
type TypeToRecord<T extends Record<string, any>, K extends Record<string, string>> = {
    [key in keyof K as K[key]]: T[K[key]];
};

/**
 *【树的自定义键配置】
 */
type TreeTypeWithouPrototype<T extends Record<string, any>, K extends Record<string, string>> = K extends {
    children: string;
}
    ? TypeToRecord<T, K>
    : TypeToRecord<T, K> & { children: TreeType<T, K> };

/**
 * 【树结构类型】
 *
 * T：树节点的类型
 *
 * K：自定义键
 *
 * 如果 K 中没有指定children，则最终返回的类型会包含children；
 *
 * 如果K中指定了children，最终返回类型将没有children，而是 K 类型中children对应值的类型。
 */
type TreeType<T extends Record<string, any>, K extends Record<string, any>> = TreeTypeWithouPrototype<T, K>[] & TreePrototypeType<TreeTypeWithouPrototype<T, K>>;

/**
 * 平铺数组转化成树
 * @param arr 平铺数组
 * @param keyMapping 使用自定义的key来作为节点数据中的id和faId
 * @returns
 */
export function treeCreate<T extends Record<string, any>, K extends TreeCreateConfig<T>["keyMapping"]>(_arr: T[], keyMapping: K): TreeType<T, K> {
    const arrData = _arr.slice();

    /** id、faId、children的key */
    const customInfo = {
        idKey: (keyMapping.id ? keyMapping.id : TreeNodeTypeKeyEnum.id) as string,
        faIdKey: (keyMapping.faId ? keyMapping.faId : TreeNodeTypeKeyEnum.faId) as string,
        childrenKey: (keyMapping.children ? keyMapping.children : TreeNodeTypeKeyEnum.children) as string,
    };

    /** 获取到数据的id值 */
    const getId = (data: T): string | null => {
        // 判断是否有keyMapping中的id，如果有
        const id: string = (data as any)[customInfo.idKey];

        if (!checkIndexType(id)) return null;
        return id;
    };

    /** 获取到数据的id值 */
    const getFaId = (data: T): string | null => {
        const faId: string = (data as any)[customInfo.faIdKey];

        if (!checkIndexType(faId)) return null;
        return faId;
    };

    /** 获取到数据的children值 */
    const getChildren = (data: T) => {
        const children = (data as any)[customInfo.childrenKey];
        if (!(children instanceof Array)) return null;
        return children;
    };

    /** 构建树*/
    const buildTree = () => {
        const tree = bindTreeFunction([]);
        const treeMap = new Map<string, T>();

        arrData.forEach((item) => {
            const id = getId(item);
            const faId = getFaId(item);
            const children = bindTreeFunction(getChildren(item) || []);

            if (id) {
                treeMap.set(id, {
                    [customInfo.idKey]: id,
                    [customInfo.faIdKey]: faId,
                    [customInfo.childrenKey]: children,
                } as unknown as T);
            }
        });

        for (let i = 0; i < arrData.length; i++) {
            const item = arrData[i];
            const faId = getFaId(item);
            const id = getId(item);

            if (!id || !faId) continue;

            const node = treeMap.get(id);
            const father = treeMap.get(faId);
            // 遍历每一个节点
            if (node) {
                if (!father) {
                    tree.push(node);
                } else {
                    const children = getChildren(father);
                    if (!(children instanceof Array)) {
                        throw new Error(`'${customInfo.childrenKey}'的值不是一个数组`);
                    }
                    children.push(node);
                }
            }
        }

        return tree;
    };

    /** 绑定原型函数 */
    const bindTreeFunction = (tree: TreeType<T, K>[]) => {
        const childrenKey = customInfo.childrenKey as keyof TreeType<T, K>;

        const prototypeFun: TreePrototypeType<TreeType<T, K>> = {
            treeForEach: (callback) =>
                treeDFS(tree, callback, {
                    childrenKey,
                }),
            treeFind: (callback) =>
                treeFind(tree, callback, {
                    childrenKey,
                }),
            treeFilter: (callback, config) =>
                treeFilter(tree, callback, {
                    ...config,
                    childrenKey,
                }),
        };

        const res = Object.create(Array.prototype, Object.getOwnPropertyDescriptors(prototypeFun));
        return Object.setPrototypeOf(tree, res);
    };

    return buildTree();
}

/**
 * 深度优先遍历
 * @param tree 树数据
 * @param callback 遍历到每个元素时的回调，参数为节点
 * @param config.childrenKey
 */
export function treeDFS<T>(tree: T[], callback?: (item: T) => void, config?: { childrenKey?: keyof T }): void {
    const childrenKey = config?.childrenKey || "children";

    /** 传递节点 获取到children的引用 */

    const getChildren = (item: T) => {
        const children = (item as any)[childrenKey];
        if (!(children instanceof Array)) {
            throw new Error("没有children字段，或在配置项参数中指定childrenKey为children的key");
        }
        return children;
    };

    const list = tree.slice();
    let current: T | undefined;

    while ((current = list.shift())) {
        callback && callback(current);

        const currentChildren = getChildren(current);

        if (currentChildren.length !== 0) {
            list.unshift(...currentChildren);
        }
    }
}

/**
 *
 * @param tree 树数据
 * @param callback 回调函数，返回值为boolean，参数接收节点类型
 * @param confi.childrenKey
 * @returns callback为true时返回的节点，如果没有符合callback的节点则返回null
 */
export function treeFind<T>(
    tree: T[],
    callback?: (item: T) => boolean,
    config?: {
        childrenKey?: keyof T;
    }
): T | null {
    const { childrenKey } = config || {};

    /** 传递节点 获取到children的引用 */
    const getChildren = (item: T) => {
        const children = (item as any)[childrenKey || "children"];
        if (!(children instanceof Array)) {
            throw new Error("没有children字段，或在配置项参数中指定childrenKey为children的key");
        }
        return children;
    };

    const list = tree.slice();
    let current: T | undefined;

    while ((current = list.shift())) {
        const children = getChildren(current);

        if (callback && callback(current)) {
            // 符合筛选callback条件的返回
            return current;
        } else {
            if (children.length !== 0) {
                list.unshift(...children);
            }
        }
    }

    return null;
}

/**
 * @param tree 树数据
 * @param callback 过滤回调函数
 * @param config.isSearch
 * 为true时，搜索过滤，搜索出节点，保留不符合条件的父节点。
 *
 * 为false时，子树过滤，不符合callback的节点以及子节点都舍弃。
 * @param config.childrenKey 树中节点用来嵌套子节点的children的key
 * @returns 节点数组
 */
export function treeFilter<T>(
    tree: T[],
    callback?: (item: T) => boolean,
    config?: {
        isSearch?: boolean;
        childrenKey?: keyof T;
    }
): T[] {
    const { childrenKey } = config || {};

    /** 传递节点 获取到children的引用 */
    const getChildren = (item: T) => {
        const children = (item as any)[childrenKey || "children"];
        if (!(children instanceof Array)) {
            throw new Error();
        }
        return children;
    };

    const { isSearch } = config || {};
    const filterTarget: T[] = [];

    const list = tree.slice();
    let current: T | undefined;

    // 未传递callback时，返回空数组
    if (!callback) return [];

    while ((current = list.shift())) {
        const children = getChildren(current);

        const filteredChildren = treeFilter(children, callback, {
            ...config,
            isSearch: true,
        });

        // 搜索过滤，只要子节点满足，则所有父级节点都保留
        if (isSearch) {
            if (filteredChildren.length > 0 || callback(current)) {
                filterTarget.push({
                    ...current,
                    [childrenKey as string]: filteredChildren,
                });
            }
        } else {
            // 子树过滤，节点不满足则舍弃所有子节点
            if (callback(current)) {
                filterTarget.push({
                    ...current,
                    [childrenKey as string]: filteredChildren,
                });
            }
        }
    }

    return filterTarget;
}

下面是使用。在 createTree 生成的节点中，会在原型上绑定 treeForeach、treeFind、treeFilter 方法，在原型上调用时会利用 ts 类型简化一些例如 childrenKey 的配置项的传递。

const nodes = [
    {
        code: "7zwqkfyqshl",
        facode: "0",
    },
    {
        code: "zzuuldvpbbc",
        facode: "7zwqkfyqshl",
    },
];
const targetTree = treeCreate(nodes, {
    id: "code",
    faId: "facode",
});

console.log("@@", targetTree);

// ------ 树遍历
let count = 0;
treeDFS(
    targetTree,
    (item) => {
        count++;
        item;
    },
    {
        childrenKey: "test",
    }
);

targetTree.treeForEach((item) => {
    count++;
    item;
});
console.log("2. 树遍历", count);

// ------ 树查找
console.log("3. 树查找");
treeFind(
    targetTree,
    (item) => {
        return item.code.includes("f");
    },
    {
        childrenKey: "test",
    }
);

const findTarget = targetTree.treeFind((item) => {
    return item.code.includes("f");
});
console.log(findTarget, findTarget?.code);

// ------ 树过滤
// 子树过滤
targetTree.treeFilter((item) => {
    return item.code.includes("f");
});

// 搜索过滤
const filterTarget = treeFilter(
    targetTree,
    (item) => {
        return item.code.includes("f");
    },
    {
        isSearch: true,
        childrenKey: "test",
    }
);
console.log("4. 树过滤", filterTarget);